import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import pandas as pd

from var_strategy.finace_manage import FinanceManage, Strategy
from var_strategy.strategies.StrategyA import StrategyA

lva = "close_var_l_ma_fib_m1"
fva = "close_var_f_ma_fib_m1"
va = "close_var_ma_fib_m1"
turn = 'turnover_ma_fib_m1'
cma = 'close_ma_fib_m1'


class BackTestV2:
    def __init__(self, data_path, fm: FinanceManage):
        self.df = self.load_data(data_path)
        self.symbols = set(self.df['symbol'])
        self.fm = fm


    def load_data(self, data_path):
        return pd.read_csv(data_path)

    def pretreatment(self):
        df = self.df
        df['date'] = df['date'].astype('datetime64[ns]')
        df['close_var_ma3_upd'] = df['close_var_ma3'] - df.groupby(['symbol'])['close_var_ma3'].shift(3)
        df['close_var_ma5_upd'] = df['close_var_ma5'] - df.groupby(['symbol'])['close_var_ma5'].shift(5)
        df['close_var_ma10_upd'] = df['close_var_ma10'] - df.groupby(['symbol'])['close_var_ma10'].shift(10)

        # df["close_ma" + "_fib_w1"] = df['close_ma5'] / df['close_ma20'] - 1.0
        df[cma] = df['close_ma20'] / df.groupby(['symbol'])['close_ma60'].shift(10) - 1.0

        df[lva + "upd"] = df[lva] - df.groupby(['symbol'])[lva].shift(5)
        df[fva + "upd"] = df[fva] - df.groupby(['symbol'])[fva].shift(5)
        df[va + "upd"] = df[va] - df.groupby(['symbol'])[va].shift(5)

        df[turn + "upd"] = df[turn] - df.groupby(['symbol'])[turn].shift(5)
        df[cma + "upd"] = df[cma] - df.groupby(['symbol'])[cma].shift(5)
        print(df["close_var_f_ma5"].quantile([.25, .5, .75]))
        print(df["close_var_f_ma20"].quantile([.25, .5, .75]))
        print(df["close_var_f_ma60"].quantile([.25, .5, .75]))
        print(df["close_var_l_ma5"].quantile([.25, .5, .75]))
        print(df["close_var_l_ma20"].quantile([.25, .5, .75]))
        print(df["close_var_l_ma60"].quantile([.25, .5, .75]))
        self.df = df.dropna(axis=0, how='any')

    def run(self):
        self.pretreatment()
        i = 0
        for symbol in self.symbols:
            print("----------" + str(i) + " :" + symbol + "-----------")
            i = i + 1
            is_buy = []
            mdata = self.df[self.df['symbol'] == symbol]
            # print(mdata["close_var_f_ma60"])
            last_row = None
            for index, row in mdata.iterrows():
                self.fm.market_day(symbol, row)
                last_row = row
            if last_row is not None:
                self.fm.pool[symbol].sell(last_row["close"], 0, "")
            # self.fm.pool[symbol].draw(mdata)
        self.stat_ret()
        print(self.fm.init_cash)

    def stat_ret(self):
        keys = self.fm.return_map.keys()
        keys = sorted(keys)
        sum_rate = 0
        for k in keys:
            print(k, self.fm.return_map[k])
            sum_rate = sum_rate + k * self.fm.return_map[k]
        count = sum(self.fm.return_map.values())
        print(count, "次交易， 平均收益率:", sum_rate / count)


if __name__ == '__main__':
    # s = Strategy()
    s = StrategyA()
    b = FinanceManage(s)
    a = BackTestV2('D:\\code\\rlearn\\stock_399344_var_model_data_v2.csv', b)
    a.run()
    # todo 下跌趋势不能买
    # todo 分步止损
    # todo 试下强化学习
    # todo 判断区间强势，价格涨幅少，但是波动小很大